import numpy as np
import matplotlib.pyplot as plt
from MUSIC_tools import generate_signal

# 参数设置
num_signals = 2
num_sensors = 4
sensor_distance = 0.5  # 传感器间距
wavelength = 1  # 波长
theta_resolution = 1  # 角度分辨率
num_samples = 1000  # 采样点数
SNR_dB = 20  # 信噪比

def root_music(array_signal, wavelength, sensor_distance):
    # root-MUSIC 算法
    M, T = array_signal.shape  # (num_sensors, num_samples)
    K = num_signals  # 信源数

    # 计算协方差矩阵
    R = np.cov(array_signal)
    # 特征值分解并取得噪声子空间
    D, U = np.linalg.eigh(R)  # D为特征值， U为特征向量
    Un = U[:, :K]
    Gn = Un @ Un.conj().T
    # 提取多项式系数并对多项式求根
    coe = np.zeros(2 * M - 1, dtype=np.complex)
    for i in range(-(M - 1), M):  # M -> sensors
        coe[-i + M-1] = np.sum(np.diag(Gn, i))
    r = np.roots(coe)  # 利用roots函数求多项式的根
    r = r[np.abs(r) < 1]  # 找出在单位圆里的根
    distances = np.abs(np.abs(r) - 1)
    sorted_indices = np.argsort(distances)
    sorted_r = r[sorted_indices][:K]
    Theta = np.arccos(np.angle(sorted_r) * wavelength / (2 * np.pi * sensor_distance))  # 计算信号到达方向角

    return Theta


# 生成仿真信号
array_signal = generate_signal(num_signals, num_sensors, sensor_distance, wavelength, SNR_dB, num_samples)

# root-MUSIC声源定向
estimated_angles = root_music(array_signal, wavelength, sensor_distance)

# 打印估计结果
print('估计结果：')
print(np.rad2deg(estimated_angles))

# 绘制声源定向结果
# plt.figure()
# plt.plot(estimated_angles, np.zeros_like(estimated_angles), 'ro')
# plt.title('声源定向结果')
# plt.xlabel('角度')
# plt.ylabel('幅度')
# plt.grid(True)
# plt.show()
